# Kestrel: An XMPP-based Many-Task Computing Scheduler
# Author: Lance Stout <lancestout@gmail.com>
#
# Copyright 2009-2010 Clemson University
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Task Executor Module

The task executor module provides access to the system shell
for running executable files.
"""

import subprocess


def execute(command):
    """Wrapper function to open a subprocess."""
    return subprocess.Popen(command,
                            shell=True,
                            stdout=subprocess.PIPE,
                            stderr=subprocess.PIPE)


class Executor(object):
    """
    Module for executing shell commands via tasks.

    Events listened for by this module:
      task/approved
      task/cancelled

    Events raised by this module:
      task/started
      task/finished
      task/cleaned
    """

    def __init__(self, kernel):
        """
        Initialize the executor for running approved task
        commands and tracking their progress.
        """
        self.kernel = kernel
        self.events = {"task/approved": self.run,
                       "task/cancelled": self.stop}
        self.kernel.register(self)
        self.process = None
        self.task = None

    def run(self, task):
        """
        Execute the given task in two stages. The first stage
        runs the executable marked as the task's command. The
        second stage is an optional cleanup phase.
        """
        self.kernel.event("task/started", task)
        self.task = task

        # Execute the task's main command and capture stdout
        # and stderr.
        try:
            self.process = execute(task.command + ' ' + str(task.task_id))
            result = self.process.wait()
            stdout, stderr = self.process.communicate()
        except:
            stdout, stderr = '', '[ERROR] Could not execute: %s' % task.command

        task.stdout, task.stderr = stdout, stderr

        # Send notification that the task's goal has been
        # completed.
        self.kernel.event("task/finished", task)

        # If the task specified a cleanup action, execute it now. If
        # the system goes offline in this phase the task will not
        # need to be resubmitted.
        if task.cleanup != '':
            try:
                self.process = execute(task.cleanup)
                result = self.process.wait()
                stdout, stderr = self.process.communicate()
            except:
                pass

        self.process = None
        self.task = None

        self.kernel.event("task/cleaned", task)

    def stop(self, data=None):
        """Terminate the currently running task."""
        if self.process is not None:
            self.process.kill()
